from typing import Dict, Any, List, Optional
import os
import sys
import importlib.util

# 定义动态导入函数
def import_module_from_path(module_name, file_path):
    """动态导入模块"""
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

# 确定当前目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 尝试导入简单向量数据库
try:
    from simple_vector_db import SimpleVectorDB, get_vector_db
except ImportError:
    # 如果失败，尝试动态导入
    vector_db_module = import_module_from_path(
        "simple_vector_db", os.path.join(current_dir, "simple_vector_db.py")
    )
    SimpleVectorDB = vector_db_module.SimpleVectorDB
    get_vector_db = vector_db_module.get_vector_db

# 导入嵌入模型
try:
    from embedding import EmbeddingModel
except ImportError:
    # 如果失败，尝试动态导入
    embedding_module = import_module_from_path(
        "embedding", os.path.join(current_dir, "embedding.py")
    )
    EmbeddingModel = embedding_module.EmbeddingModel


# Initialize the vector database and embedding model as module-level singletons
# This ensures we only create one instance even if the tool is called multiple times
_vector_db = None
_embedding_model = None


def get_embedding_model() -> EmbeddingModel:
    """Get or create the embedding model singleton."""
    global _embedding_model
    if _embedding_model is None:
        _embedding_model = EmbeddingModel()
    return _embedding_model


def knowledge_retrieval_tool(query: str, tool_context=None) -> Dict[str, Any]:
    """
    从企业知识库中检索与用户查询相关的文档片段。

    Args:
        query: 用户提出的问题或查询。
        tool_context: ADK 工具上下文，可用于访问会话状态或认证信息。

    Returns:
        一个包含检索到的文档片段列表的字典，
        例如: {"status": "success", "documents": ["文档片段1...", "文档片段2..."]}
             {"status": "not_found", "message": "未找到相关文档。"}
    """
    print(f"--- 工具：knowledge_retrieval_tool 被调用，查询={query} ---")
    
    try:
        # 获取向量数据库和嵌入模型实例
        vector_db = get_vector_db()
        embedding_model = get_embedding_model()
        
        # 将用户查询转换为向量表示
        query_embedding = embedding_model.embed(query)
        
        # 在向量数据库中搜索相似文档
        search_results = vector_db.search(
            query_vectors=query_embedding,
            limit=3  # 返回前3个最相关的结果
        )
        
        # 检查是否有结果
        if not search_results:
            return {
                "status": "not_found",
                "message": "在企业知识库中未找到与您查询直接相关的文档。"
            }
        
        # 提取文档文本并返回
        documents = [result["text"] for result in search_results]
        
        return {
            "status": "success",
            "documents": documents
        }
        
    except Exception as e:
        print(f"知识检索工具出错: {e}")
        return {
            "status": "error",
            "error_message": f"检索文档时发生错误: {str(e)}"
        }


def initialize_knowledge_base_with_examples() -> None:
    """用示例数据初始化知识库，仅用于演示目的。"""
    vector_db = get_vector_db()
    embedding_model = get_embedding_model()
    
    # 示例文档
    documents = [
        {
            "text": "HR政策文档摘录1：公司年假政策规定每位员工每年可享有20天带薪假期。新员工需要工作满3个月后方可申请休假。休假申请需提前至少一周提交给直属经理审批。",
            "metadata": {"source": "HR政策手册", "category": "假期政策", "updated": "2023-09-01"}
        },
        {
            "text": "HR政策文档摘录2：病假规定超过3天需要提供医生证明。病假期间，员工薪资将按照公司规定的比例发放。紧急情况下，请通过电话或短信及时通知您的直属经理。",
            "metadata": {"source": "HR政策手册", "category": "假期政策", "updated": "2023-09-01"}
        },
        {
            "text": "技术设置指南摘录1：配置公司VPN，请按照以下步骤操作：1.从内部门户下载VPN客户端 2.使用您的公司邮箱和密码登录 3.选择最近的服务器位置 4.连接后，您将可以访问内部资源。",
            "metadata": {"source": "IT设置指南", "category": "网络配置", "updated": "2023-10-15"}
        },
        {
            "text": "技术设置指南摘录2：安装新软件前，请确保您的操作系统已更新到最新版本。所有软件安装必须通过公司软件中心进行，不得自行下载安装未经批准的软件。",
            "metadata": {"source": "IT设置指南", "category": "软件安装", "updated": "2023-10-15"}
        },
        {
            "text": "新员工入职指南：您的第一天将安排参观办公室，了解公司文化，并与您的团队成员见面。请携带您的身份证件和银行账户信息，以便HR部门完成入职手续。",
            "metadata": {"source": "入职手册", "category": "新员工", "updated": "2023-08-10"}
        },
        {
            "text": "远程工作政策：公司允许员工每周远程工作不超过2天。远程工作期间，员工需要保持在线状态并参加所有必要的会议。每周必须在办公室工作至少3天。",
            "metadata": {"source": "工作场所政策", "category": "工作安排", "updated": "2023-11-05"}
        }
    ]
    
    # 为每个文档生成向量嵌入
    document_texts = [doc["text"] for doc in documents]
    embeddings = embedding_model.embed(document_texts)
    
    # 准备要插入的数据
    data_to_insert = []
    for i, (doc, embedding) in enumerate(zip(documents, embeddings)):
        data_to_insert.append({
            "id": i,
            "vector": embedding,
            "text": doc["text"],
            "metadata": doc["metadata"]
        })
    
    # 删除现有集合并创建新的
    vector_db.drop_collection()
    # 插入文档
    vector_db.add_documents(data_to_insert)
    
    print(f"成功初始化知识库，添加了 {len(data_to_insert)} 个文档。") 